import os
import re
import argparse
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
import torch

def process_file(file_path):
    if not os.path.isfile(file_path):
        print(f"file path: {file_path} is not valid")
        return []

    mar_line_pattern = re.compile(r'\d+\s+MAR')
    ofl_line_pattern = re.compile(r'\d+\s+OFL')

    line_counts = []
    found_first_mar = False
    count = 0
    with open(file_path, 'r') as file:
        data_started = False
        for line in file:
            if len(line_counts) == 398:
                x = 1

            if not data_started and line.startswith("record#"):
                data_started = True
                continue

            if not data_started:
                continue

            if mar_line_pattern.match(line):
                if found_first_mar:
                    if count > 0:
                        line_counts.append(count)
                    count = 1  # Reset count for the new MAR line
                else:
                    found_first_mar = True
                    count = 1  # Start counting from this MAR line
            elif found_first_mar and not ofl_line_pattern.match(line):
                count += 1

        # If the file ends without another MAR, ignore the last count
    if found_first_mar and count > 0:
        # Only add the count if we have found at least one MAR entry and there's a non-zero count.
        average_count = sum (line_counts) / len(line_counts)
        line_counts.append(average_count)
        


    return line_counts

def DGI_reconstruction(y, patterns, batch_size, num_patterns, img_h, img_w, g_factors=0.5):
    patterns = torch.reshape(patterns, [num_patterns, img_h * img_w])
    y = torch.reshape(y, [batch_size, num_patterns, 1])

    compute1 = torch.transpose(
        (patterns - torch.matmul(torch.ones([num_patterns, 1]), torch.mean(patterns, dim=0).unsqueeze(0))), dim0=0,
        dim1=1)
    compute2 = torch.sum(patterns, dim=1)

    gammer = g_factors * torch.mean(y, dim=(1, 2)) / torch.mean(compute2)

    temp = (gammer.unsqueeze(-1) * compute2).unsqueeze(-1)

    # DGI porcess
    DGI = torch.matmul(compute1, y - temp).squeeze(-1)
    DGI = torch.reshape(DGI, [batch_size, 1, img_h, img_w])
    mean_DGI = torch.mean(DGI)
    std_DGI = torch.std(DGI)
    DGI = (DGI - mean_DGI) / std_DGI

    return DGI

parser = argparse.ArgumentParser(description="Process a list of files and generate plots.")
parser.add_argument('files', nargs='+', help='List of file paths to process')
parser.add_argument('--output-path', type=str, required=True, help='Path where the image will be saved')
parser.add_argument('--output-name', type=str, required=True, help='Name of the output image file (without extension)')
parser.add_argument('--pt-file-path', type=str, required=True, help='the pt file path')

args = parser.parse_args()
file_list = args.files
output_path = args.output_path
output_name = args.output_name
ptFilePath = args.pt_file_path

print(f"Received file list from command line: {file_list}")
print(f"Output image will be saved at: {output_path}/{output_name}.png")

total_photo_num = []
for file_path in file_list:
    counts = process_file(file_path)
    total_photo_num.extend(counts)

if not os.path.exists(output_path):
    os.makedirs(output_path)



y_np = total_photo_num  # 这里用具体的实际数据代替
plt.figure(figsize=(15, 6))

# 转换为PyTorch张量
y_torch = torch.tensor(y_np, dtype=torch.float32)

y_torch = y_torch.reshape([1, 4096, 1, 1])

patterns_path = ptFilePath
patterns = torch.load(patterns_path, map_location=torch.device('cpu'))

DGI = DGI_reconstruction(y_torch, patterns, 1, patterns.shape[0], patterns.shape[2], patterns.shape[3], 0.5)

print(DGI.shape)
DGI = DGI.squeeze().numpy()
plt.imshow(DGI,cmap="gray")  # 显示灰度图像
plt.title('Reconstructed Image using DGI')
plt.axis('off')
output_image_path = Path(output_path) / Path(output_name + ".png")
plt.savefig(output_image_path, bbox_inches='tight')
plt.close()